
#read "encoder.mpl":

# This function builds a graph
# the number of a variable vertex is the corresponding column number
# of H. the number of a check vertex is the row number + total number
# of columns of H. The function returns a record with three fields.
# 1. Variable nodes;
# 2. Check nodes;
# 3. The constructed Tanner Graph.
TannerGraph := proc(H)
local rownum, colnum, E, i, j, V, C;
    (rownum, colnum) := LinearAlgebra:-Dimension(H);
    E := {};
    V := {};
    C := {};
    for i from 1 to rownum do
        for j from 1 to colnum do
            if H(i, j) = 1 then
                V := {op(V), j};
                C := {op(C), colnum + i};
                E := {op(E), {j, colnum + i}};
            end if;
        od;
    od;
    [V, C, GraphTheory:-Graph(E)]
end proc;

# a testing case for building the Tanner graph
#L := TannerGraph(GH[1]):
#GraphTheory:-DrawGraph(L[3]);

# examine whether all the values of the parity check nodes are 0.
parityChecks := proc(V, vSize, cSize)
local i;
    for i from 1 to cSize do
        if V[vSize + i] = 1 then
            return false;
        end if;
    od;
    true;
end proc;

# The iterative decoder for BEC. The decoder takes in a parity check
# matrix H, an iteration limit lim, and a vector of received code bits
# S of length n. It returns a vector of the corrected code bits.  For
# now, we require the length of S equals n.
decodeBEC := proc(S, H, lim)
local triple, messageBuffers, vNodes, cNodes, tGraph, vSize, cSize, v, c, c2;
local toRemoveV, neighborv, iter;
    triple := TannerGraph(H);
    (vNodes, cNodes) := (triple[1], triple[2]);
    (vSize, cSize)   := (nops(vNodes), nops(cNodes));
    tGraph := triple[3];  # the Tanner's graph of H

    # 1. initialization
    messageBuffers := [op(S), seq(0, i = 1..cSize)];

    # start iteration.
    iter := 0;
    while (iter < lim) do
        # 2. Direct recovery
        toRemoveV := {};
        for v in vNodes do
            if messageBuffers[v] <> e then
                for c in GraphTheory:-Neighbors(tGraph, v) do
                    messageBuffers[c] := (messageBuffers[c] + messageBuffers[v]) mod 2;
                od;
                # remove the variable node v and all its adjacent edges
                tGraph := GraphTheory:-DeleteVertex(tGraph, [v]);
                # mark v to be removed from vNodes
                toRemoveV := {op(toRemoveV), v};
            end if;

        od;
        # update variable node set.
        vNodes := vNodes minus toRemoveV;

        # 3. Substitution recovery
        for c in cNodes do
            if GraphTheory:-Degree(tGraph, c) = 1 then
                neighborv := GraphTheory:-Neighbors(tGraph, c)[1];
                # do the substitution
                messageBuffers[neighborv] := messageBuffers[c];
                for c2 in GraphTheory:-Neighbors(tGraph, neighborv) do
                    messageBuffers[c2] := (messageBuffers[c2] + messageBuffers[neighborv]) mod 2;
                od;
                # remove the variable node v and all its adjacent edges
                tGraph := GraphTheory:-DeleteVertex(tGraph, [neighborv]);
                # remove neighborv from vNodes
                vNodes := vNodes minus {neighborv};
            end if;

        od;
        # if parity checks, return.
        if parityChecks(messageBuffers, vSize, cSize) then
            return messageBuffers;
        end if;

        # otherwise, this is terminated according to the iteration limit.
        iter := iter + 1;
    od;
    print("Warning: terminate due to iteration limit!");
    messageBuffers;
end proc;

# for testing decoding
# H := Matrix([[1, 1, 1, 1, 0, 0], [0, 0, 1, 1, 0, 1], [1, 0, 0, 1, 1, 0]]);
# S := [e, 0, e, 0, 1, 1]:
# lim := 5:
# decodeBEC(S, H, lim);

# This function computes the initial value for q_ij(0)
# which is 1 - Pr_i := 1 - Pr_i(ci = 1 | yi)
# given the channel error rate P. We assume the probability of
# 0 -> 1 and 1 -> 0 are both equivalent to P.
initialProb := proc(P, yi)
    if yi = 0 then
        return (1 - P);
    end if;
    P;
end proc;


# compute response from check node c_j.
# i.e., r_ji(0)
computeCheckNodeResponse := proc(qmat, i, j, neighbors)
    local V, prod, k;
    # variable nodes except v_i
    V := convert(neighbors, set) minus {i};
    prod := 1;
    for k in V do
        prod := prod * (1 - 2 * (1 - qmat(k, j)))
    od;
    0.5 * (prod + 1);
end proc;


# compute part of the response from variable node v_i.
multiplyResponse1 := proc(rmat, i, j,  neighbors, vSize, P, yi)
local V, prod, k, j2;
    V := convert(neighbors, set) minus {(j+vSize)};
    prod := initialProb(P, yi);
    for k in V do
        j2 := (k - vSize);
        prod := prod * rmat(j2, i);
    od;
    prod;
end proc;

multiplyResponse2 := proc(rmat, i, j,  neighbors, vSize, P, yi)
local V, prod, k, j2;
    V := convert(neighbors, set) minus {(j+vSize)};
    prod := 1 - initialProb(P, yi);
    for k in V do
        j2 := (k - vSize);
        prod := prod * (1 - rmat(j2, i));
    od;
    prod;
end proc;

# This function implements the LDPC decoder using the sum product
# algorithm in probability domain.
# S: received codeword bits.
# H: parity-check matrix .
# P: the error/erasure probability. e.g. error rate of a BSC channel.
decodeSPA := proc(S, H, lim, P)
local triple, vNodes, cNodes, tGraph, vSize, cSize, i, j, v, c, qij, rji, iter;
local qmat, rmat; #qmat and rmat's elements are q_ij(0) and r_ij(0),
                  #respectively;
local neighbors; # the set of neighbors of a node.
local qij0, qij1, K; # two symbolic expressions for solving Kij;
local C; # stores the estimation of the transmitted codewords
         # determined at step 4.
local coef0, coef1, prod0, prod1, HC;

    triple := TannerGraph(H);
    (vNodes, cNodes) := (triple[1], triple[2]);
    (vSize, cSize)   := (nops(vNodes), nops(cNodes));
    tGraph := triple[3];  # the Tanner's graph of H

    # initialization of the vector containing estimated codewords that
    # is transmitted
    C := [seq(0, i = 1..vSize)];

    # element q_ij is the message sent from variable node v_i to check
    # node c_j.
    qmat := Matrix(vSize, cSize, 0);
    # element r_ji is the message sent from check node c_i to variable
    # node v_i.
    rmat := Matrix(cSize, vSize, 0);

    # 1. initialization of q_ij(0)
    # q_ij(1)s are implied by 1 - q_ij(0).
    for v in vNodes do
        i := v;
        # yi := S[i]
        qij := initialProb(P, S[i]);
        for c in GraphTheory:-Neighbors(tGraph, v) do
            j := c - vSize;
            qmat(i, j) := qij;
        od;
    od;

    iter := 0;
    while (iter < lim) do
        # 2. computes r_ji(0), where check node c_j and variable node v_i
        # are neighbors.
        for c in cNodes do
            j := c - vSize;
            neighbors := GraphTheory:-Neighbors(tGraph, c);
            for v in neighbors do
                i := v;
                rji := computeCheckNodeResponse(qmat, i, j, neighbors);
                rmat(j, i) := rji;
            od;
        od;

        # 3. variable nodes update the response messages q_ij to check
        # nodes via equation solving.
        for v in vNodes do
            i := v;
            neighbors := GraphTheory:-Neighbors(tGraph, v);
            for c in neighbors do
                j := c - vSize;
                coef0 := multiplyResponse1(rmat, i, j,  neighbors, vSize, P, S[i]);
                qij0 := Kij * coef0;
                coef1 := multiplyResponse2(rmat, i, j,  neighbors, vSize, P, S[i]);
                qij1 := Kij * coef1;
                K := solve(qij0 + qij1 = 1, Kij);
                #update q_ij(0).
                qmat(i, j) := K * coef0;
            od;
        od;

        # 4. compute the estimation of codewords.
        for v in vNodes do
            i := v;
            neighbors := GraphTheory:-Neighbors(tGraph, v);
            prod0 := initialProb(P, S[i]);
            prod1 := 1 - prod0;
            for c in neighbors do
                j := c - vSize;
                prod0 := prod0 * rmat(j, i);
                prod1 := prod1 * (1 - rmat(j, i));
            od;
            # FIXME: computation of K seems to be not necessary.
            K := solve(Ki * (prod0 + prod1) = 1, Ki);
            if (K * prod1) > 0.5 then
                C := subsop(i=1, C); # C[i] := 1;
            elif (K * prod1) <= 0.5 then
                C := subsop(i=0, C); # C[i] := 0;
            else
                print(K*prod1);
                print("overflow may happen!");
                error "terminate due to error";
            end if;
        od;

        # termination condition via checking parity check matrix
        HC := LinearAlgebra[Modular]:-Multiply(2, H, Vector(C));
        if (add(i, i = HC) = 0) then
            print("Decoding is successful!");
            return [true, iter+1, C];
        end if;
#        myprint("current iteration:", iter);
        iter := iter + 1;
    od;
    print("Warning: decoding terminated due to iteration limit!");
    [false, iter, C]
end proc;

# for testing SPA-decoder
# H := Matrix([[1, 1, 1, 1, 0, 0], [0, 0, 1, 1, 0, 1], [1, 0, 0, 1, 1, 0]]);
# S := [1, 0, 1, 0, 0, 1]: # the first bit is wrong
# lim := 5:
# P := 1/6;
# C := decodeSPA(S, H, lim, P);



# computes the initial message used in step 1 of the belief
# propagation decoder.
initialMessageBSC := proc(yi, P)
    ((-1)^yi) * log((1-P)/P)
end proc:

# given a returns the sign of a (1 for +, -1 for -) and the absolute
# value |a|
signAndMagnitude := proc(Lqij)
    #print(Lqij);
    if evalf(Lqij) >= 0 then
        return [1, abs(Lqij)];
    end if;
    [-1, abs(Lqij)];
   #[sign(evalf(Lqij)), abs(Lqij)]
end proc:

# the phi function
phi := proc(x)
    if x = 0 then
        print("x is:");
        print(x);
    end if;
    log((exp(x) + 1)/(exp(x) - 1))
end proc;

# L(c_i)
lci := proc(yi, P)
    if yi = 0 then
        return log((1-P)/P);
    end if;
    log(P/(1-P))
end proc;


# computes the response from check node c_i to variable node v_i.
# i.e. L(rji)
computeCheckNodeResponseBP := proc(qmat, i, j, neighbors)
    local V, prod, s, k, k2, Lqij, aij, bij, pair;
    # variable nodes except v_i
    V := convert(neighbors, set) minus {i};

    # this is because row weight = 1 from the back parity-check matrix
    # FIXME: may introduce bug
    if nops(V) = 0 then
        return 0;
    end if;

    prod := 1;
    for k in V do
        Lqij := qmat(k, j);
        aij := signAndMagnitude(Lqij)[1];
        prod := prod * aij;
    od;

    s := 0;
    for k2 in V do
        Lqij := qmat(k2, j);
        bij := signAndMagnitude(Lqij)[2];
        s := s + phi(bij);
    od;

    # temporary fix, if s = 0, then we tell variable node we don't
    # know.
    # FIXME: may introduce bug
    if s = 0 then
        error "s = 0 !!!";
        return 0;
    end if;

    prod * phi(s)
end proc:


# This function implements the LDPC decoder using the sum product
# algorithm in probability domain.
# S: received codeword bits.
# H: parity-check matrix .
# P: the error/erasure probability. e.g. error rate of a BSC channel.
decodeBP := proc(S, H, lim, P)
local triple, vNodes, cNodes, tGraph, vSize, cSize, i, j, j2, v, c, c2, iter;
local qmat, rmat; #qmat and rmat's elements are L(q_ij) and L(r_ij), respectively.
local neighbors; # the set of neighbors of a node.
local C; # stores the estimation of the transmitted codewords
         # determined at step 4.
local HC;
local Lqij, Lrji, LQi;

    triple := TannerGraph(H);
    (vNodes, cNodes) := (triple[1], triple[2]);
    (vSize, cSize)   := (nops(vNodes), nops(cNodes));
    tGraph := triple[3];  # the Tanner's graph of H

    # initialization of the vector containing estimated codewords that
    # is transmitted
    C := [seq(0, i = 1..vSize)];

    # element L(q_ij) is the message sent from variable node v_i to check
    # node c_j.
    qmat := Matrix(vSize, cSize, 0);
    # element L(r_ji) is the message sent from check node c_i to variable
    # node v_i.
    rmat := Matrix(cSize, vSize, 0);

    # 1. initialization of L(q_ij) is log(q_ij(1)/q_ij(0)).
    for v in vNodes do
        i := v;
        # yi := S[i]
        Lqij := initialMessageBSC(S[i], P);
        for c in GraphTheory:-Neighbors(tGraph, v) do
            j := c - vSize;
            qmat(i, j) := Lqij;
        od;
    od;

    iter := 0;
    while (iter < lim) do
        # 2. computes r_ji(0), where check node c_j and variable node v_i
        # are neighbors.
        for c in cNodes do
            j := c - vSize;
            neighbors := GraphTheory:-Neighbors(tGraph, c);
            for v in neighbors do
                i := v;
                Lrji := computeCheckNodeResponseBP(qmat, i, j, neighbors);
                rmat(j, i) := Lrji;
            od;
        od;


        # 3. variable nodes update the response messages q_ij to check
        # nodes via equation solving.
        for v in vNodes do
            i := v;
            neighbors := GraphTheory:-Neighbors(tGraph, v);
            for c in neighbors do
                j := c - vSize;
                Lqij := lci(S[i], P);
                for c2 in (convert(neighbors, set) minus {c}) do
                    j2 := c2 - vSize;
                    Lqij := Lqij + rmat(j2, i);
                od;
                #update L(q_ij).
                qmat(i, j) := Lqij;
            od;
        od;


        # 4. compute the estimation of codewords.
        for v in vNodes do
            i := v;
            LQi := lci(S[i], P);
            neighbors := GraphTheory:-Neighbors(tGraph, v);
            for c in neighbors do
                j := c - vSize;
                LQi := LQi + rmat(j, i);
            od;
            if (evalf(LQi) < 0) then
                C := subsop(i=1, C); # C[i] := 1;
            elif (evalf(LQi) >= 0) then
                C := subsop(i=0, C); # C[i] := 0;
            else
                print(LQi);
                #print(evalf(LQi));
                #print("matrix rmat");
                #print(rmat);
                #print(convert(rmat,matrix));
                error "overflow or underflow happens";
            end if;
        od;
#        print("current iteration:");
#        print(iter);

        # termination condition via checking parity check matrix
        HC := LinearAlgebra[Modular]:-Multiply(2, H, Vector(C));
        if (add(i, i = HC) = 0) then
            print("belief propagation decoding is successful!");
            return [true, iter+1, C];
        end if;
        iter := iter + 1;
    od;
    print("Warning: decoding terminated due to iteration limit!");
    [false, iter, C];
end proc;


# main entry for LDPC decoding on BSC:
# option = 1 : in log-domain (belief propagation)
# option = 2 : in probability-domain (sum product)
decodeBSC := proc(S, H, lim, P, opt)
    if opt = 1 then
        return decodeBP(S, H, lim, P);
    elif opt = 2 then
        return decodeSPA(S, H, lim, P);
    else
        print("User error: unrecognized option value");
    end if;
end proc;

# for testing LDPC-decoder on BSC
# H := Matrix([[1, 1, 1, 1, 0, 0], [0, 0, 1, 1, 0, 1], [1, 0, 0, 1, 1, 0]]);
# S := [1, 0, 1, 0, 0, 1]: # the first bit is wrong
# lim := 2:
#P := 1/6;
# BP
#C1 := decodeBSC(S, H, lim, P, 1);
# SPA
#C2 := decodeBSC(S, H, lim, P, 2);
